from clickhouse_driver import Client
from typing import Union
import pandas as pd
import re
from  datetime import datetime,date,time,timedelta
# import Config


# conf_ = Config()
# config = conf_.get_conf
host = "8.142.23.85"
port = "9000"
user = "default"
pwd = "cqzp123456"


class ClickHouse:

    def __init__(self):
        # host = config.get('ClickHouse', 'host')
        # port = config.getint('ClickHouse', 'port')
        # user = config.get('ClickHouse', 'user')
        # pwd = config.get('ClickHouse', 'password')
        self.client = Client(host=host,port=port,user=user ,password=pwd)

    def read_sql(self,sql):
        data, columns = self.client.execute(sql, columnar=True, with_column_types=True)
        df = pd.DataFrame({re.sub(r'\W', '_', col[0]): d for d, col in zip(data, columns)})
        return df

    def get_type_dict(self,tb_name,db_name):
        sql = f"select name, type from system.columns where database='{db_name}' and table='{tb_name}';"
        df = self.read_sql(sql)
        df = df.set_index('name')
        type_dict = df.to_dict('dict')['type']
        return type_dict

    def to_sql(self,df, tb_name,db):
        type_dict = self.get_type_dict(tb_name,db)
        columns = list(type_dict.keys())
        # 类型处理
        for i in range(len(columns)):
            col_name = columns[i]
            col_type = type_dict[col_name]
            if 'Date' in col_type:
                df[col_name] = pd.to_datetime(df[col_name])
            if 'Int' in col_type:
                df[col_name] = df[col_name].astype('int')
            elif 'Float' in col_type:
                df[col_name] = df[col_name].astype('float')
            elif col_type == 'String':
                df[col_name] = df[col_name].astype('str').fillna('')
        # df数据存入clickhouse
        cols = ','.join(columns)
        data = df.to_dict('records')
        self.client.execute(f"INSERT INTO {db}.{tb_name} ({cols}) VALUES", data, types_check=True)

    def get_col_names(self, db: str):
        sql = "SELECT name FROM system.tables WHERE database = '"+db+"'"
        col_ls = self.client.execute(sql)
        return col_ls
    

    def insert_dataframe(self, df: pd.DataFrame, db: str, collection: str):
        self.to_sql(df,collection,db)

    def read_dataframe(self, db: str, collection: str, filter: Union[dict, None] = None):
        type_dict = self.get_type_dict(collection,db)
        columns = list(type_dict.keys())
        if filter is not None:
            f_sql = "1=1"
            for key in filter:
                f_sql += " and " + key + " "+filter[key]
            sql = f"select * from {db}.{collection} where {f_sql}"
            # print(sql)
            res_data = pd.DataFrame(self.client.execute(f"select * from {db}.{collection} where {f_sql}"),columns=columns)
            return res_data
        else:
            res_data = pd.DataFrame(self.client.execute(f"select * from {db}.{collection} where 1=1"),columns=columns)
            return res_data

    def get_md_data_from_clickhouse(self, db: str, collection: str,symbol: str,start_date:datetime,end_date:datetime):
        return self.client.execute(f"select * from {db}.{collection} where 1=1 and symbol = '{symbol}' and trading_date >= '{datetime.datetime.strptime(start_date,'%Y-%m-%d %H:%M:%S')}' and trading_date <= '{datetime.datetime.strptime(end_date,'%Y-%m-%d %H:%M:%S')}'")

    def execute_sql(self,sql:str):
        return self.client.execute(sql)
#调用实例
# ck = ClickHouse()
# f_dict = {"contract":"='AG2106'"} #筛选条件
# ck.read_dataframe("zp_futures","cn_md_bars_day",f_dict)

# df = pd.read_csv("/Users/weihe/Documents/quantification_data/symbol_and_exchange.csv",index_col=0)
# ck.insert_dataframe(df,"zp_futures","cn_md_symbol_and_exchange")